from os import path

import torch
import torch.nn as nn
import numpy as np

import config
from model import vgg16_101
from pretrain import build_data

def eval_model(net, model_path=None, mode='val', batch_size=None):
    '''
        check the accuracy of given model.
        
        `net`: the model object
        `model_path`: state_dict path to be load, if given
        `mode`: decide the dataset and whether top5 acc will be checked
    '''

    if mode == 'val':
        sample_sum = config.VAL_NUM
    elif mode == 'test':
        sample_sum = config.TEST_NUM
    else:
        raise ValueError('Evaluation mode should be "val" or "test".')

    # load state dict, if given
    if model_path is not None and path.exists(model_path):
        net.load_state_dict(torch.load(model_path))

    net.eval()  # convert to eval mode
    net = net.cuda()  # move to gpu
    if batch_size is None:
        batch_size = config.BATCH_SIZE  # default batch size
    _, valloader = build_data(data_set='val', batch_size=batch_size)  # get val dataset

    with torch.no_grad():

        top1_num, top5_num =  0., 0.  # init stats

        for data in valloader:
            
            # go through the model
            inputs, targets = data
            inputs, targets = inputs.cuda(), targets.cuda()
            batch_num = inputs.size()[0]
            outputs = net(inputs)

            # count top-1 accuracy
            _, predicted = torch.max(outputs.data, 1)
            top1_num += (targets == predicted).sum()

            # when not training, count top-5 accuracy
            if mode == 'test':
                _, top5 = outputs.topk(5, dim=1)
                top5_num += (targets == top5).sum()
        
        top1_acc = top1_num / sample_sum
        top5_acc = top5_num / sample_sum
        if mode == 'val':
            return top1_acc
        else:
            return top1_acc, top5_acc

if __name__ == '__main__':
    net = vgg16_101()
    model_path = 'output/model-vgg16.pth'
    acc, acc_top5 = eval_model(net, model_path, mode='test', batch_size=1)
    print('Top1 accuracy on test set: %.2f%%' % (100 * acc))
    print('Top5 accuracy on test set: %.2f%%' % (100 * acc_top5))
